class TreeNode(object):
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right


class Solution(object):
    def deleteNode(self, root, key):
        """
        :type root: TreeNode
        :type key: int
        :rtype: TreeNode
        """
        if not root:
            return None

        def search(node):
            if not node:
                return None
            if node.val == key:
                if node.left:
                    left_node = node.left
                    if not left_node.right:
                        left_node.right = node.right
                        return left_node
                    else:
                        right_node = left_node
                        while right_node.right.right:
                            right_node = right_node.right
                        out_node = right_node.right
                        right_node.right = None
                        out_left = out_node
                        while out_left.left:
                            out_left = out_left.left
                        out_left.left = left_node
                        out_node.right = node.right
                        return out_node
                else:
                    return node.right
            else:
                if key < node.val:
                    node.left = search(node.left)
                if key > node.val:
                    node.right = search(node.right)
                return node

        return search(root)
